import torch
from torch.nn import functional as F
from torch import nn
from pytorch_lightning.core.lightning import LightningModule
import pytorch_lightning as pl

import torch.optim as optim
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt

import pickle

if torch.cuda.is_available():
    device = torch.device("cuda") 
else:
    device = torch.device("cpu")

import json
import argparse
import sys

import matplotlib.pyplot as plt
import pandas as pd
from IPython.core.display import HTML
import sys

sys.path = sys.path[:8]
sys.path

['/home/jupyter/nwp-downscale/notebooks',
 '/opt/conda/envs/ilan/lib/python39.zip',
 '/opt/conda/envs/ilan/lib/python3.9',
 '/opt/conda/envs/ilan/lib/python3.9/lib-dynload',
 '',
 '/opt/conda/envs/ilan/lib/python3.9/site-packages',
 '/opt/conda/envs/ilan/lib/python3.9/site-packages/IPython/extensions',
 '/home/jupyter/.ipython']

Load config and results

config_path = '../experiments/eval_leingan_ens10_tp_stacked_input.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_ens10_tp_stacked_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_ens10_tp_stacked_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))
leingan_ens10_tp_stacked_preds = plt.imread(folder+'sample_predictions.png', format='png')

config_path = '../experiments/eval_leingan_single_forecast_tp_input.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_single_forecast_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_single_forecast_preds = plt.imread(folder+'sample_predictions.png', format='png')
leingan_single_forecast_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))

config_path = '../experiments/eval_wgan-gp_single_forecast_tp.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# wgan_gp_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
wgan_gp_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))
wgan_gp_preds = plt.imread(folder+'sample_predictions.png', format='png')


config_path = '../experiments/eval_wgan-gp-smoothed_single_forecast_tp.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# wgan_gp_smoothed_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
wgan_gp_smoothed_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))
wgan_gp_smoothed_preds = plt.imread(folder+'sample_predictions.png', format='png')

config_path = '../experiments/eval_leingan_ens10_tp_random_input.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_ens10_tp_random_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_ens10_tp_random_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))
leingan_ens10_tp_random_preds = plt.imread(folder+'sample_predictions.png', format='png')

config_path = '../experiments/eval_leingan_single_forecast_tp_with_consts_input.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_single_forecast_with_consts_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_single_forecast_with_consts_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))
leingan_single_forecast_with_consts_preds = plt.imread(folder+'sample_predictions.png', format='png')

config_path = '../experiments/eval_leingan_ens10_tp_random_added_vars_TCW-T-CAPE-CIN.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_ens10_added_vars_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_ens10_added_vars_preds = plt.imread(folder+'sample_predictions.png', format='png')
leingan_ens10_added_vars_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))

config_path = '../experiments/eval_leingan_single_forecast_tp_added_vars_TCW.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_ens10_added_vars_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_single_forecast_added_TCW_preds = plt.imread(folder+'sample_predictions.png', format='png')
leingan_single_forecast_added_TCW_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))

config_path = '../experiments/eval_leingan_ens10_tp_random_input_ens_mean_L1_weighted_gen_loss.json'
args = json.load(open(config_path))
parser = argparse.ArgumentParser(args)
parser.set_defaults(**args)
args, _ = parser.parse_known_args()
folder = args.save_hparams["save_dir"] + args.save_hparams["run_name"] + str(args.save_hparams["run_number"])+"/"
# leingan_ens10_added_vars_full_field_metrics = pickle.load(open(folder+"full_field_eval_metrics.pkl", "rb"))
leingan_ens10_tp_random_mean_l1_reg_preds = plt.imread(folder+'sample_predictions.png', format='png')
leingan_ens10_tp_random_mean_l1_reg_patch_metrics = pickle.load(open(folder+"patch_eval_metrics.pkl", "rb"))

scalar_metrics = pd.DataFrame(columns=["model", "crps", "avg_pool_crps", "max_pool_crps", "rmse"])

models = [
#           ('leingan_single_forecast_full_field', leingan_single_forecast_full_field_metrics, leingan_single_forecast_preds), 
#           ('leingan_ens10_tp_stacked_full_field', leingan_ens10_tp_stacked_full_field_metrics, leingan_ens10_tp_stacked_preds, 
#           ('wgan_gp_full_field', wgan_gp_full_field_metrics, wgan_gp_preds),
#           ('wgan_gp_smoothed_full_field', wgan_gp_smoothed_full_field_metrics, wgan_gp_smoothed_preds), 
#           ('leingan_single_forecast_with_consts_full_field', leingan_single_forecast_with_consts_preds), 
#           ('leingan_ens10_tp_random_full_field', leingan_ens10_tp_random_full_field_metrics, leingan_ens10_tp_random_preds, 
#           ('leingan_ens10_added_vars_full_field', leingan_ens10_added_vars_full_field_metrics, leingan_ens10_added_vars_preds),
          ('leingan_single_forecast_patch', leingan_single_forecast_patch_metrics, leingan_single_forecast_preds), 
          ('leingan_ens10_tp_stacked_patch', leingan_ens10_tp_stacked_patch_metrics, leingan_ens10_tp_stacked_preds), 
#           ('wgan_gp_patch', wgan_gp_patch_metrics, wgan_gp_preds),
#           ('wgan_gp_smoothed_patch', wgan_gp_smoothed_patch_metrics, wgan_gp_smoothed_preds), 
          ('leingan_single_forecast_with_consts_patch', leingan_single_forecast_with_consts_patch_metrics, leingan_single_forecast_with_consts_preds), 
          ('leingan_ens10_tp_random_patch', leingan_ens10_tp_random_patch_metrics, leingan_ens10_tp_random_preds), 
#           ('leingan_ens10_added_vars_patch', leingan_ens10_added_vars_patch_metrics, leingan_ens10_added_vars_preds), 
          ('leingan_single_forecast_added_TCW_patch', leingan_single_forecast_added_TCW_patch_metrics, leingan_single_forecast_added_TCW_preds), 
          ('leingan_ens10_tp_mean_l1_reg_patch', leingan_ens10_tp_mean_l1_reg_patch_metrics, leingan_ens10_tp_mean_l1_reg_preds), 
        ]

for (model_name, metrics, preds) in models: 
    scalar_metrics = scalar_metrics.append({"model": model_name, "crps":metrics["crps"], 
                  "avg_pool_crps":metrics["avg_pool_crps"], 
                  "max_pool_crps":metrics["max_pool_crps"], 
                  "rmse":metrics["rmse"]#, "fss": metrics["fss"]
                                           }, ignore_index=True)
scalar_metrics

	model	crps	avg_pool_crps	max_pool_crps	rmse
0	leingan_single_forecast_patch	0.351096	0.343050	0.586708	1.056507
1	leingan_ens10_tp_stacked_patch	0.386971	0.384850	0.695095	1.344685
2	leingan_single_forecast_with_consts_patch	0.357854	0.350497	0.597780	1.120847
3	leingan_ens10_tp_random_patch	0.347917	0.340552	0.588700	1.016366
4	leingan_single_forecast_added_TCW_patch	0.367192	0.360441	0.617958	1.064971
5	leingan_ens10_tp_mean_l1_reg_patch	0.337140	0.328909	0.568120	0.989351

Rank Histograms

plt.figure(figsize=(20,10))
for (model_name, metrics, preds) in models: 
    plt.plot(metrics['rankhist']/np.sum(metrics['rankhist']), label=model_name, linewidth=2)
plt.hlines(y = 1/31, xmin=0, xmax=31, color='k', linestyle='--', linewidth=2)
plt.xlabel('rank')
plt.ylabel('normalised frequency')
plt.title('Rank Histogram')
plt.legend()

<matplotlib.legend.Legend at 0x7f01e17ab880>

plt.figure(figsize=(20,10))
for (model_name, metrics, preds) in models: 
    plt.plot(np.abs(metrics['rankhist']/np.sum(metrics['rankhist']) - 1/31), label=model_name, linewidth=2)
# plt.hlines(y = 1/31, xmin=0, xmax=31, color='k', linestyle='--', linewidth=2)
plt.xlabel('rank')
plt.ylabel('normalised frequency')
plt.title('Rank Histogram deviation from uniform')
plt.legend()

<matplotlib.legend.Legend at 0x7f01e17ba070>

Reliability diagrams

plt.figure(figsize=(20,10))

for (model_name, metrics, preds) in models: 
    relative_freq, forecast_probs, samples = metrics['reliability_4']
    plt.plot(forecast_probs, relative_freq, label=model_name)
plt.plot(np.linspace(0, 0.95, len(forecast_probs)), np.linspace(0, 0.95, len(forecast_probs)), 'k--')
plt.xlabel('Forecast Probability')
plt.ylabel('Observed Frequency')
plt.title('Reliability Diagram (>4)')
plt.legend()

<matplotlib.legend.Legend at 0x7f01e1c730a0>

plt.figure(figsize=(20,10))

for (model_name, metrics, preds) in models: 
    relative_freq, forecast_probs, samples = metrics['reliability_1']
    plt.plot(forecast_probs, relative_freq, label=model_name)
plt.plot(np.linspace(0, 0.95, len(forecast_probs)), np.linspace(0, 0.95, len(forecast_probs)), 'k--')
plt.xlabel('Forecast Probability')
plt.ylabel('Observed Frequency')
plt.title('Reliability Diagram (>1)')
plt.legend()

<matplotlib.legend.Legend at 0x7f01e153b0a0>

TP histograms

for (model_name, metrics, preds) in models:
    plt.figure()
    preds_hist_x = (np.array(metrics['preds_hist'][1][:-1]) + np.array(metrics['preds_hist'][1][1:]))/2
    plt.semilogy(preds_hist_x, metrics['preds_hist'][0]/sum(metrics['preds_hist'][0]), label = f'Predictions')
    true_hist_x = (np.array(metrics['true_hist'][1][:-1]) + np.array(metrics['true_hist'][1][1:]))/2
    plt.semilogy(true_hist_x, metrics['true_hist'][0]/sum(metrics['true_hist'][0]), label = 'Observations')
    plt.title(model_name)

plt.figure()
for (model_name, metrics, preds) in models:
    preds_hist_x = (np.array(metrics['preds_hist'][1][:-1]) + np.array(metrics['preds_hist'][1][1:]))/2
    plt.plot(preds_hist_x[:50], (metrics['preds_hist'][0] - metrics['true_hist'][0])[:50], label = model_name)
plt.legend()

<matplotlib.legend.Legend at 0x7f01e0cf4760>

array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True])

TP Mean

plt.figure(figsize=(20,10))
plt.plot(models[0][1]['true_mean'], 1, 'o', label='Observations mean')

for (model_name, metrics, preds) in models:
    plt.plot(metrics['preds_mean'], 1, 'o', label=f'{model_name} mean')
    
plt.legend()

<matplotlib.legend.Legend at 0x7f01e116afd0>

Visualising predictions samples

for (model_name, metrics, preds) in models:
    plt.figure(figsize=(20, 45))
    plt.imshow(preds)
    plt.title(model_name)

!jupyter nbconvert view_eval.ipynb --to html

[NbConvertApp] Converting notebook view_eval.ipynb to html
[NbConvertApp] Writing 5604588 bytes to view_eval.html